﻿using Org.BouncyCastle.Crypto;
using Org.BouncyCastle.Crypto.Digests;
using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Crypto.Signers;
using Org.BouncyCastle.Math;
using Org.BouncyCastle.Math.EC;
using Org.BouncyCastle.Math.EC.Multiplier;
using Org.BouncyCastle.Security;
using Org.BouncyCastle.Utilities;
using System;

namespace BuddHa.Common.SMUtil.Implementation.V2
{
    internal class SM2Signer
    {
        private readonly IDsaKCalculator _kCalculator = new RandomDsaKCalculator();
        private readonly IDigest _digest;
        private readonly IDsaEncoding _encoding;

        private ECDomainParameters _ecParams;
        private ECPoint _pubPoint;
        private ECKeyParameters _ecKey;
        private byte[] _z;
        private bool _generatingSignature;
        public SM2Signer() : this(StandardDsaEncoding.Instance, new SM3Digest())
        {
            _generatingSignature = false;
        }

        public SM2Signer(IDsaEncoding encoding, IDigest digest)
        {
            _encoding = encoding;
            _digest = digest;
        }

        public virtual void Init(ICipherParameters parameters, byte[] userId )
        {
            ICipherParameters baseParam;
            if (parameters is ParametersWithID parametersWithID)
            {
                baseParam = parametersWithID.Parameters;
                userId = parametersWithID.GetID();

                if (userId.Length >= 8192)
                    throw new ArgumentException("SM2 UserID必须小于2^16 bits long");
            }
            else
            {
                baseParam = parameters;
            }

            if (baseParam is ParametersWithRandom parametersWithRandom)
            {
                _ecKey = (ECKeyParameters)parametersWithRandom.Parameters;
                _ecParams = _ecKey.Parameters;
                _kCalculator.Init(_ecParams.N, parametersWithRandom.Random);
            }
            else
            {
                _ecKey = (ECKeyParameters)baseParam;
                _ecParams = _ecKey.Parameters;
                _kCalculator.Init(_ecParams.N, new SecureRandom());
            }
            _pubPoint = new FixedPointCombMultiplier().Multiply(_ecParams.G, ((ECPrivateKeyParameters)_ecKey).D).Normalize();
            _z = GetZ(userId);
        }

        public byte[] GenerateSignature(byte[] data)
        {
            if (_generatingSignature)
            {
                throw new CryptoException("当前组件正在签名");
            }
            _generatingSignature = true;
            Reset();
            _digest.BlockUpdate(data, 0, data.Length);
            byte[] eHash = DigestUtilities.DoFinal(_digest);

            BigInteger n = _ecParams.N;
            BigInteger e = CalculateE(n, eHash);
            BigInteger d = ((ECPrivateKeyParameters)_ecKey).D;

            BigInteger r, s;

            ECMultiplier basePointMultiplier = new FixedPointCombMultiplier();

            // 5.2.1 Draft RFC:  SM2 Public Key Algorithms
            do // generate s
            {
                BigInteger k;
                do // generate r
                {
                    // A3
                    //k = new BigInteger("15386121922915615803613132160314876142244421174293470239507300753591497434485");
                    k = _kCalculator.NextK();

                    // A4
                    ECPoint p = basePointMultiplier.Multiply(_ecParams.G, k).Normalize();

                    // A5
                    r = e.Add(p.AffineXCoord.ToBigInteger()).Mod(n);
                }
                while (r.SignValue == 0 || r.Add(k).Equals(n));

                // A6
                BigInteger dPlus1ModN = BigIntegers.ModOddInverse(n, d.Add(BigIntegers.One));

                s = k.Subtract(r.Multiply(d)).Mod(n);
                s = dPlus1ModN.Multiply(s).Mod(n);
            }
            while (s.SignValue == 0);

            // A7
            try
            {
                return _encoding.Encode(_ecParams.N, r, s);
            }
            catch (Exception ex)
            {

                throw new CryptoException("无法编码签名内容: " + ex.Message, ex);
            }
            finally
            {
                _generatingSignature = false;
            }
        }

        public virtual bool VerifySignature(byte[] signature)
        {
            try
            {
                BigInteger[] rs = _encoding.Decode(_ecParams.N, signature);

                return VerifySignature(rs[0], rs[1]);
            }
            catch (Exception)
            {
            }

            return false;
        }

        public virtual void Reset()
        {
            if (_z != null)
            {
                _digest.Reset();
                _digest.BlockUpdate(_z, 0, _z.Length);
            }
        }

        private byte[] GetZ(byte[] userID)
        {
            AddUserID(_digest, userID);

            AddFieldElement(_digest, _ecParams.Curve.A);
            AddFieldElement(_digest, _ecParams.Curve.B);
            AddFieldElement(_digest, _ecParams.G.AffineXCoord);
            AddFieldElement(_digest, _ecParams.G.AffineYCoord);
            AddFieldElement(_digest, _pubPoint.AffineXCoord);
            AddFieldElement(_digest, _pubPoint.AffineYCoord);

            return DigestUtilities.DoFinal(_digest);
        }

        private void AddUserID(IDigest digest, byte[] userID)
        {
            int len = userID.Length * 8;
            digest.Update((byte)(len >> 8));
            digest.Update((byte)len);
            digest.BlockUpdate(userID, 0, userID.Length);
        }

        private void AddFieldElement(IDigest digest, ECFieldElement v)
        {
            byte[] p = v.GetEncoded();
            digest.BlockUpdate(p, 0, p.Length);
        }

        private BigInteger CalculateE(BigInteger n, byte[] message)
        {
            // TODO Should hashes larger than the order be truncated as with ECDSA?
            return new BigInteger(1, message);
        }

        private bool VerifySignature(BigInteger r, BigInteger s)
        {
            BigInteger n = _ecParams.N;

            // 5.3.1 Draft RFC:  SM2 Public Key Algorithms
            // B1
            if (r.CompareTo(BigInteger.One) < 0 || r.CompareTo(n) >= 0)
                return false;

            // B2
            if (s.CompareTo(BigInteger.One) < 0 || s.CompareTo(n) >= 0)
                return false;

            // B3
            byte[] eHash = DigestUtilities.DoFinal(_digest);

            // B4
            BigInteger e = CalculateE(n, eHash);

            // B5
            BigInteger t = r.Add(s).Mod(n);
            if (t.SignValue == 0)
                return false;

            // B6
            ECPoint q = ((ECPublicKeyParameters)_ecKey).Q;
            ECPoint x1y1 = ECAlgorithms.SumOfTwoMultiplies(_ecParams.G, s, q, t).Normalize();
            if (x1y1.IsInfinity)
                return false;

            // B7
            return r.Equals(e.Add(x1y1.AffineXCoord.ToBigInteger()).Mod(n));
        }
    }
}
